package com.lyf.boot.sse;

import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

/**
 * sse服务，可用于上传进度进度监听等
 * @author Luyufan
 * @date 2023/12/12 17:08
 **/
@Slf4j
public class SSEServer {

    private static final Map<Object, SseEmitter> SSE_EMITTER_MAP = new ConcurrentHashMap<>();

    public static SseEmitter connect(Object connectId) {
        try {
            SseEmitter sseEmitter = new SseEmitter();
            sseEmitter.onCompletion(completionCallBack(connectId));
            sseEmitter.onError(errorCallBack(connectId));
            sseEmitter.onTimeout(timeoutCallBack(connectId));
            SSE_EMITTER_MAP.put(connectId, sseEmitter);
            return sseEmitter;
        } catch (Exception e) {
            log.error("{}连接异常", connectId, e);
        }
        return null;
    }

    /**
     * 单条消息发送 {{@link SseEmitter#send(Object)}}
     * @param connectId 连接Id
     * @param message 消息
     */
    public static void sendMessage(Object connectId, Object message) {
        sendMessage(connectId, message, null);
    }

    /**
     * 单挑消息发送
     * @param connectId 连接id
     * @param message 消息
     * @param mediaType 媒体类型
     */
    public static void sendMessage(Object connectId, Object message, MediaType mediaType) {
        if(SSE_EMITTER_MAP.containsKey(connectId)){
            try {
                SSE_EMITTER_MAP.get(connectId).send(message, mediaType);
            } catch (IOException e) {
                log.error("{}发送异常", connectId, e);
            }
        }
    }

    /**
     * 批量发送
     * @param message 消息
     */
    public static void batchSendMessage(Object message){
        batchSendMessage(message, null);
    }

    /**
     *
     * @param message 消息
     * @param mediaType 媒体类型
     */
    public static void batchSendMessage(Object message, MediaType mediaType){
        SSE_EMITTER_MAP.forEach((k, v) -> {
            try {
                v.send(message, mediaType);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        });
    }


    private static Runnable completionCallBack(Object connectId){
        return () -> {
            log.info("连接：{}结束连接", connectId);
            remove(connectId);
        };
    }

    private static Consumer<Throwable> errorCallBack(Object connectId){
        return throwable -> {
            log.info("连接：{}异常", connectId, throwable);
            remove(connectId);
        };
    }

    private static Runnable timeoutCallBack(Object connectId){
        return () -> {
            log.info("连接：{}超时", connectId);
            remove(connectId);
        };
    }

    public static void remove(Object connectId){
        SSE_EMITTER_MAP.remove(connectId);
        log.info("连接数为{}", SSE_EMITTER_MAP.size());
    }

}
